/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#ifndef __ORDINAL_HH__
#define __ORDINAL_HH__

#include <cstring>
#include <vector>

#include "bogackishampine.hh"
#include "fastonebigheader.hh"
#include "model.hh"
#include "pianovec.hh"
#include "rating.hh"

namespace playerpiano
{
  class Ordinal
    {
      private:
        unsigned int n_workers;
        unsigned int n_labels;
        unsigned int n_items;
        float* tau;
        float* gamma;
        float* loglambda;
        float* chi;
        float prior_sq_stddev;
        unsigned int failcount;
        float sqrt_n_labels;

        static inline float
        square (float x)
          {
            return x * x;
          }

        template<typename Expr>
        static void
        dlogpratingdtau (float*        dtau,
                         float&        dlogbeta,
                         const Rating& rating,
                         unsigned int  true_label,
                         float         beta,
                         const float*  tauptr,
                         const Expr&   e,
                         float         loglambda,
                         unsigned int  n_labels)
          {
            float p[n_labels];
            float betalambda = beta * fastexp (loglambda + e[0]);
            float max;

            p[0] = 0;
            max = p[0];

            for (unsigned int i = 0; i + 1 < n_labels; ++i)
              {
                float v = tauptr[i] + e[i + 1] + i + 0.5;

                p[i + 1] = p[i] + betalambda * (true_label - v);

                if (p[i + 1] > max)
                  {
                    max = p[i + 1];
                  }
              }

            float denominator = 0;
            float ddenominator_dlogbeta = 0;

            dlogbeta = p[rating.label] - max;

            for (unsigned int i = 0; i < n_labels; ++i)
              {
                float d = p[i] - max;
                p[i] = fastexp (d);
                denominator += p[i];
                ddenominator_dlogbeta += p[i] * d;
              }

            dlogbeta -= ddenominator_dlogbeta / denominator;

            float cdf = 0;

            for (unsigned int i = n_labels - 1; i > 0; --i)
              {
                cdf += p[i] / denominator;
                dtau[i - 1] = cdf;

                if (i <= rating.label)
                  {
                    dtau[i - 1] -= 1;
                  }

                dtau[i - 1] *= betalambda;
              }
          }

      public:
        Ordinal (float*       _tau,
                 float*       _gamma,
                 float*       _loglambda,
                 float*       _chi,
                 float        prior_stddev,
                 unsigned int _n_workers,
                 unsigned int _n_labels,
                 unsigned int _n_items)
          : n_workers (_n_workers),
            n_labels (_n_labels),
            n_items (_n_items),
            tau (_tau),
            gamma (_gamma),
            loglambda (_loglambda),
            chi (_chi),
            prior_sq_stddev (square (prior_stddev)),
            failcount (0),
            sqrt_n_labels (sqrt (n_labels))
          {
          }

        Ordinal (Model* model)
          : n_workers (1U << strtoul (model->lookup<char*> ("--n_worker_bits").first, 0, 0)),
            n_labels (strtoul (model->lookup<char*> ("--n_labels").first, 0, 0)),
            n_items (strtoul (model->lookup<char*> ("--n_items").first, 0, 0)),
            tau (model->lookup<float*> ("tau").first),
            gamma (model->lookup<float*> ("gamma").first),
            loglambda (model->lookup<float*> ("loglambda").first),
            chi (model->lookup<float*> ("chi").first),
            prior_sq_stddev (1.0),
            failcount (0),
            sqrt_n_labels (sqrt (n_labels))
          {
          }

        void
        apply_deferred_prior (unsigned int  w,
                              uint64_t      s,
                              uint64_t      t,
                              float         eta,
                              float         rho)
          {
            if (t > s)
              {
                float powt = fastpow (t, 1.0f - rho);
                float sigma = prior_sq_stddev;
                float pows = fastpow (s, 1.0f - rho);
                float decay = fastexp (eta * (pows - powt) / 
                                       (n_items * (1.0f - rho) * sigma));

                float* tauptr = tau + w * (n_labels - 1);

                for (unsigned int k = 0; k + 1 < n_labels; ++k)
                  { 
                    float dtautmp = (1.0 - decay) * (gamma[k] - tauptr[k]);

                    tauptr[k] += dtautmp;
                    gamma[k] += dtautmp / (1.0 + n_workers);
                  }

                /* loglambda and chi */
      
                float dloglambdatmp = (1.0 - decay) * (*chi - loglambda[w]);
      
                loglambda[w] += dloglambdatmp;
                *chi += dloglambdatmp / (1.0 + n_workers);
              }
          }

        class DFunc
          {
            private:
              const Rating&             rating;
              const std::vector<float>& pz;
              float                     beta;
              const float*              tauptr;
              float                     loglambda;
              unsigned int              n_labels;

            public:
              DFunc (const Rating&             _rating,
                     const std::vector<float>& _pz,
                     float                     _beta,
                     const float*              _tauptr,
                     float                     _loglambda,
                     unsigned int              _n_labels)
                  : rating (_rating),
                    pz (_pz),
                    beta (_beta),
                    tauptr (_tauptr),
                    loglambda (_loglambda),
                    n_labels (_n_labels)
                {
                }

              template<typename Expr>
              PianoVec 
              operator() (const float&,
                          const Expr& e) const
                {
                  PianoVec rv (n_labels, 0);

                  for (unsigned int j = 0; j < n_labels; ++j)
                    {
                      float tmpdtau[n_labels - 1];
                      float dlogbeta;

                      dlogpratingdtau (tmpdtau,
                                       dlogbeta,
                                       rating,
                                       j,
                                       beta,
                                       tauptr,
                                       e,
                                       loglambda,
                                       n_labels);

                      rv[0] += pz[j] * dlogbeta;

                      for (unsigned int k = 0; k + 1 < n_labels; ++k)
                        {
                          rv[k + 1] += pz[j] * tmpdtau[k];
                        }
                    }

                  return rv;
                }
          };

        void
        update (uint64_t                  t,
                float                     eta,
                float                     rho,
                float                     beta,
                const std::vector<float>& pz,
                const Rating&             rating)
          {
            float* tauptr = tau + rating.worker * (n_labels - 1);
            DFunc dfunc (rating,
                         pz,
                         beta,
                         tauptr,
                         loglambda[rating.worker],
                         n_labels);
            float etat = eta * fastpow (t, -rho);
            PianoVec s (n_labels, 0);

            std::pair<bool, PianoVec> step = 
              bogacki_shampine_2_3 (dfunc,
                                    0.0f,
                                    s,
                                    etat,
                                    5e-3f * sqrt_n_labels,
                                    etat / 1000.0,
                                    etat,
                                    1000);

            if (! step.first)
              {
                ++failcount;
              }
            else
              {
                loglambda[rating.worker] += step.second[0];
                *chi += step.second[0] / (1.0 + n_workers);
    
                for (unsigned int k = 0; k + 1 < n_labels; ++k)
                  {
                    tauptr[k] += step.second[k + 1];
                    gamma[k] += step.second[k + 1] / (1.0 + n_workers);
                  }
              }
          }

        float
        logprating (const Rating&   rating,
                    unsigned int    true_label,
                    float           beta) const
          {
            float* tauptr = tau + rating.worker * (n_labels - 1);
            float betalambda = beta * fastexp (loglambda[rating.worker]);
            float logfactors[n_labels];
            unsigned int argmax;

            logfactors[0] = 0;
            argmax = 0;

            for (unsigned int i = 0; i + 1 < n_labels; ++i)
              {
                float v = tauptr[i] + i + 0.5;

                logfactors[i + 1] = 
                  logfactors[i] + betalambda * (true_label - v);

                if (logfactors[i + 1] > logfactors[argmax])
                  {
                    argmax = i + 1;
                  }
              }

            float denominator = 0;

            for (unsigned int i = 0; i < n_labels; ++i)
              {
                float d = logfactors[i] - logfactors[argmax];
                float expd = fastexp (d);
                denominator += expd;
              }

            return logfactors[rating.label] 
                   - logfactors[argmax]
                   - fastlog (denominator);
          }

        float
        logprating (const Rating&   rating,
                    unsigned int    true_label,
                    float           beta,
                    float&          dlogbeta,
                    float&          ddlogbeta) const
          {
            float* tauptr = tau + rating.worker * (n_labels - 1);
            float betalambda = beta * fastexp (loglambda[rating.worker]);
            float logfactors[n_labels];
            unsigned int argmax;

            logfactors[0] = 0;
            argmax = 0;

            for (unsigned int i = 0; i + 1 < n_labels; ++i)
              {
                float v = tauptr[i] + i + 0.5;

                logfactors[i + 1] = 
                  logfactors[i] + betalambda * (true_label - v);

                if (logfactors[i + 1] > logfactors[argmax])
                  {
                    argmax = i + 1;
                  }
              }

            float denominator = 0;
            float ddenominator_dlogbeta = 0;
            float ddenominator_ddlogbeta = 0;

            for (unsigned int i = 0; i < n_labels; ++i)
              {
                float d = logfactors[i] - logfactors[argmax];
                float expd = fastexp (d);
                denominator += expd;
                ddenominator_dlogbeta += expd * d;
                ddenominator_ddlogbeta += expd * d * (d + 1.0);
              }

            dlogbeta = logfactors[rating.label] 
                       - logfactors[argmax] 
                       - ddenominator_dlogbeta / denominator;

            ddlogbeta = logfactors[rating.label] 
                        - logfactors[argmax] 
                        + square (ddenominator_dlogbeta / denominator)
                        - ddenominator_ddlogbeta / denominator;

            return logfactors[rating.label] 
                   - logfactors[argmax]
                   - fastlog (denominator);
          }

        static void
        ensure (Model* model)
          {
            unsigned int n_workers (1U << strtoul (model->lookup<char*> ("--n_worker_bits").first, 0, 0));
            unsigned int n_labels (strtoul (model->lookup<char*> ("--n_labels").first, 0, 0));
            model->ensure ("tau",
                           n_workers * (n_labels - 1) * sizeof (float));
            model->ensure ("gamma", (n_labels - 1) * sizeof (float));
            model->ensure ("loglambda", n_workers * sizeof (float));
            model->ensure ("chi", sizeof (float));
          }

        unsigned int
        failures () const
          {
            return failcount;
          }
    };
}

#endif // __ORDINAL_HH__
